#include "mapping/sliding_window.h"

namespace mapping {

SlidingWindow::SlidingWindow(const Eigen::Matrix3d& intrinsic_matrix, const Eigen::Matrix4d& extrinsic_matrix, const std::shared_ptr<Publisher> publisher_ptr) : extrinsic_matrix_(extrinsic_matrix), publisher_ptr_(publisher_ptr) {
    is_init_ = false;
    imu_init_bias_ = Eigen::Vector3d::Zero();

    memset(&state_, 0.0, sizeof(State));

    last_imu_timestamp_ = 0;

    last_wheel_pose_ = SE2(0,0,0);

    camera_id_ = 0;
    frame_id_ = 0;
    feature_id_ = 0;

    feature_database_ptr_ = std::make_shared<FeatureDatabase>();
    feature_matcher_ptr_ = common::make_unique<FeatureMatcher>(feature_database_ptr_);
    triangulate_ptr_ = common::make_unique<Triangulate>(feature_database_ptr_, intrinsic_matrix, extrinsic_matrix, camera_id_);
    
    frame_ptr_ = std::make_shared<Frame>(feature_database_ptr_, intrinsic_matrix, extrinsic_matrix, camera_id_);
    sliding_window_optimizer_ptr_ = common::make_unique<SlidingWindowOptimizer>(frame_ptr_, intrinsic_matrix, extrinsic_matrix, camera_id_);
    
    is_using_imu_factor_ = sliding_window_optimizer_ptr_->is_using_imu_factor_;
    is_using_wheel_factor_ = sliding_window_optimizer_ptr_->is_using_wheel_factor_;
    is_using_visual_factor_ = sliding_window_optimizer_ptr_->is_using_visual_factor_;
}

SlidingWindow::~SlidingWindow() {

}

void SlidingWindow::StartMapping(const std::deque<ImuData>& imu_data, const std::deque<WheelData>& wheel_data, const std::vector<FeaturePoint>& feature_data) {
    if(!is_init_) {
      Initialize(imu_data, wheel_data, feature_data);
      if(publisher_ptr_ != nullptr) {
          publisher_ptr_->SetState(Eigen::Vector3d(state_.p[0], state_.p[1], 0.0));
          // publisher_ptr_->SetState(Eigen::Vector3d(wheel_state_.p[0], wheel_state_.p[1], 0.0));
      }
    } else {
      StartProcessSliddingWindow(imu_data, wheel_data, feature_data);
      if(publisher_ptr_ != nullptr) {
          publisher_ptr_->SetState(Eigen::Vector3d(state_.p[0], state_.p[1], 0.0));
          // publisher_ptr_->SetState(Eigen::Vector3d(wheel_state_.p[0], wheel_state_.p[1], 0.0));
      }
    }
}

void SlidingWindow::Initialize(const std::deque<ImuData>& imu_data, const std::deque<WheelData>& wheel_data, const std::vector<FeaturePoint>& feature_data) {
    if(wheel_data.empty() || imu_data.empty()) {
        return;
    }
    Eigen::Vector3d imu_init_bias = Eigen::Vector3d::Zero();

    // LOG(INFO) << " Imu data: " << imu_data.size() << ", " << wheel_data.size() << std::endl;
    is_init_ = ImuBiasEstimate(imu_init_bias, wheel_data, imu_data);
    if(!is_init_) {
      return;
    }

    imu_init_bias_ = imu_init_bias;

    SE2 current_wheel_pose = SE2(wheel_data.back().x, wheel_data.back().y, wheel_data.back().t);

    state_.p = current_wheel_pose.translation();
    state_.v[0] = wheel_data.back().vehicle_speed;
    state_.v[1] = 0;
    state_.R = current_wheel_pose.so2();
    state_.ba = imu_init_bias.head<2>();
    state_.bg = imu_init_bias[2];
    state_.pose = SE2::matrix4d(state_.R, state_.p);

    // initial imu
    imu_options_.init_ba_ = state_.ba;
    imu_options_.init_bg_ = state_.bg;
    imu_options_.noise_gyro_ = 1.5;
    imu_options_.noise_acce_ = 10.0;
    imu_preintegration_ptr_ = std::make_shared<ImuPreIntegration>(imu_options_);
        
    // initial wheel 
    wheel_options_.sigma_x = 0.001;
    wheel_options_.sigma_y = 0.001;
    wheel_options_.sigma_t = 100.0;
    wheel_preintegration_ptr_ = std::make_shared<WheelPreIntegration>(wheel_options_);

    last_imu_ = TransformImuData(imu_data.back());

    LOG(INFO) << "state_ " << state_.p.transpose() << ", " << state_.R.log() << std::endl;
    LOG(INFO) << "current_wheel_pose " << current_wheel_pose.translation().transpose() << ", " << current_wheel_pose.so2().log() << std::endl;

    last_wheel_pose_ = current_wheel_pose;  

    last_imu_timestamp_ = imu_data.back().timestamp;

    // update feature database
    for(const feature::FeaturePoint& feat_point : feature_data) {
      feature_database_ptr_->UpdateFeatureDatabase(camera_id_, frame_id_, feature_id_, feat_point);
      ++feature_id_;
    }

    Eigen::Vector2d vel_odom = Eigen::Vector2d(wheel_data.back().vehicle_speed, 0.0);
    Eigen::Vector2d vel_world = state_.R.matrix() * vel_odom;
    frame_ptr_->InsertSlidingWindowData(frame_id_, state_, imu_preintegration_ptr_, wheel_preintegration_ptr_, vel_world);

    frame_id_ += 1;

    imu_data_.clear();

    wheel_state_ = state_;
}

bool SlidingWindow::ImuBiasEstimate(Eigen::Vector3d& imu_bias, const std::deque<WheelData>& wheel_data, const std::deque<ImuData>& imu_data) {
    bool is_init = false;
    if(wheel_data.back().vehicle_speed > 1e-8) {
      return is_init;
    }
    for(const auto& imu : imu_data) {
        imu_data_.push_back(imu);
    }

    if(imu_data_.size() >= 200) {
       for(const auto& imu : imu_data_) {
          imu_bias[0] += imu.ax;
          imu_bias[1] += imu.ay;
          imu_bias[2] += imu.gz;
       }

       imu_bias[0] /= imu_data_.size();
       imu_bias[1] /= imu_data_.size();
       imu_bias[2] /= imu_data_.size();

       std::cout << "imu_bias ===== : " << imu_bias.transpose() << std::endl;

       is_init = true;
    }
    return is_init;
}

Eigen::Vector3d SlidingWindow::TransformImuData(const ImuData& imu_data) {
  Eigen::Vector3d imu_vector(imu_data.ax, imu_data.ay, imu_data.gz);
  return imu_vector;
}

Eigen::Vector3d SlidingWindow::TransformWheelData(const WheelData& wheel_data) {
  Eigen::Vector3d wheel_pose(wheel_data.x, wheel_data.y, wheel_data.t);
  return wheel_pose;
}

void SlidingWindow::StartProcessSliddingWindow(const std::deque<ImuData>& imu_data, const std::deque<WheelData>& wheel_data, const std::vector<FeaturePoint>& feature_data) {
    // imu preintegration
    for(const auto& imu : imu_data) {
        Eigen::Vector3d current_imu = TransformImuData(imu);
        float delta_time = imu.timestamp - last_imu_timestamp_;
        if((fabs(wheel_data.back().vehicle_speed) < 1e-8) && (fabs(wheel_data.back().vehicle_speed) < 1e-8)) {
            last_imu_timestamp_ = imu.timestamp;
            last_imu_ = current_imu;
            continue;
        }
        imu_preintegration_ptr_->Integrate(last_imu_, current_imu, delta_time);
        last_imu_timestamp_ = imu.timestamp;
        last_imu_ = current_imu;

        // LOG(INFO) << "imu.timestamp ============== : " << imu.timestamp << std::endl;

    }
    // wheel preintegration
    SE2 last_wheel_pose;
    uint32_t idx = 0;
    for(const auto& wheel : wheel_data) {
        SE2 current_wheel_pose = SE2(wheel.x, wheel.y, wheel.t);
        if(idx == 0) {
            last_wheel_pose = current_wheel_pose;
        }

        wheel_preintegration_ptr_->Integrate(last_wheel_pose, current_wheel_pose);

        // LOG(INFO) << "last_wheel_pose " << last_wheel_pose.translation().transpose() << std::endl;

        last_wheel_pose = current_wheel_pose;

        // LOG(INFO) << "current_wheel_pose state_ " << current_wheel_pose.translation().transpose() << std::endl;
        
        // LOG(INFO) << "wheel.timestamp ============== : " << wheel.timestamp << std::endl;

        ++idx;
    }
    
    if((last_wheel_pose.translation() - last_wheel_pose_.translation()).norm() < 0.5) {
        return;
    }

    // LOG(INFO) << "last_wheel_pose state_ " << last_wheel_pose.translation().transpose() << std::endl;
    // LOG(INFO) << "last_wheel_pose_ " << last_wheel_pose_.translation().transpose() << std::endl;
    // LOG(INFO) << "wheel_preintegration_ptr_ " << wheel_preintegration_ptr_->dp_.transpose() << std::endl;

    State last_state = state_;

    LOG(INFO) << "last state_ " << state_.p.transpose() << ", " << state_.R.log() << std::endl;
    state_ = imu_preintegration_ptr_->Predict(last_state);

    wheel_preintegration_ptr_->Predict(wheel_state_);

    if(!is_using_imu_factor_) {
        state_ = wheel_state_;
    }

    LOG(INFO) << "predict state_ " << state_.p.transpose() << ", " << state_.R.log() << std::endl;
    // LOG(INFO) << "current_wheel_pose " << current_wheel_pose.translation().transpose() << ", " << current_wheel_pose.so2().log() << std::endl;
    
    LOG(INFO) << "wheel_state_ === : " << wheel_state_.p.transpose() << ", " << wheel_state_.R.log() << std::endl;
    // LOG(INFO) << "wheel_data vehicle_speed : " << wheel_data.back().vehicle_speed << std::endl;

    SliddingWindow(imu_data.back(), wheel_data.back(), feature_data);
    
    imu_preintegration_ptr_ = std::make_shared<ImuPreIntegration>(imu_options_);
    wheel_preintegration_ptr_ = std::make_shared<WheelPreIntegration>(wheel_options_);

    last_wheel_pose_ = last_wheel_pose;

    frame_id_ += 1;

}

void SlidingWindow::SliddingWindow(const ImuData& imu_data, const WheelData& wheel_data, const std::vector<FeaturePoint>& feature_data) {
    Eigen::Vector2d vel_odom = Eigen::Vector2d(wheel_data.vehicle_speed, 0.0);
    Eigen::Vector2d vel_world = state_.R.matrix() * vel_odom;

    LOG(INFO) << "vel_world : " << vel_world.transpose() << std::endl;

    frame_ptr_->InsertSlidingWindowData(frame_id_, state_, imu_preintegration_ptr_, wheel_preintegration_ptr_, vel_world);

    uint64_t last_frame_id = frame_id_-1;
    feature_matcher_ptr_->DescriptorMatch(feature_id_, camera_id_, frame_id_, last_frame_id, feature_data);

    // feature_matcher_ptr_->RobustDescriptorMatch(feature_id_, camera_id_, frame_id_, last_frame_id, feature_data);

    FeatureTriangulate();
    
    // LOG(INFO) << " frame_id_ ===== : " << frame_id_ << std::endl;

    Optimizer();
    
    // LOG(FATAL) << " frame_id_ ===== : " << frame_id_ << std::endl;
}

void SlidingWindow::FeatureTriangulate() {
    std::map<uint64_t, Eigen::Matrix4d> poses;
    for(uint32_t idx = 0; idx < frame_ptr_->frame_ids_queue_.size(); ++idx) {
        const State state = frame_ptr_->state_queue_.at(idx);
        poses[frame_ptr_->frame_ids_queue_.at(idx)] = SE2::matrix4d(state.R, state.p);

        // LOG(INFO) << " SE2::matrix4d(state.R, state.p) ===== : " << SE2::matrix4d(state.R, state.p) << std::endl;
    }

    const std::vector<Vector3d> world_point = triangulate_ptr_->TriangulateDLT(poses);
    
    if(publisher_ptr_ != nullptr) {
        publisher_ptr_->SetCurrentLandmarks(world_point);
    }

    // LOG(INFO) << " world_point ===== : " << world_point.size() << std::endl;

    // LOG(FATAL) << " sliding windows is full " << std::endl;

}

void SlidingWindow::Optimizer() {

    sliding_window_optimizer_ptr_->Solve();
    
    frame_ptr_->MarginalizationOldestFrame(sliding_window_optimizer_ptr_->is_using_visual_factor_);

    state_ = frame_ptr_->state_queue_.back();

    LOG(INFO) << "current state_ : " << state_.p.transpose() << ", " << state_.R.log() << ", v : " << state_.v.transpose() << ", bg : " << state_.bg << ", ba : " << state_.ba.transpose() << std::endl;

}

} // end mapping

